import torch
import torch.nn.functional as F


import torch.nn as nn

from models.component.normlization import get_nonspade_norm_layer
from models.networks.basenetwork import BaseNetwork
from option import Options


class MultiscaleDiscriminator(BaseNetwork):
    def __init__(self, opt):
        super(MultiscaleDiscriminator, self).__init__()
        self.opt = opt
        for i in range(opt.num_D):
            subnet_D = NLayerDiscriminator(opt)
            self.add_module('discriminator_%d' % i, subnet_D)

    def downsample(self, input):
        return F.avg_pool2d(input, kernel_size=3,
                            stride=2, padding=[1, 1],
                            count_include_pad=False)

    def forward(self, input):
        result = []
        get_intermediate_features = not self.opt.no_ganFeat_loss
        for name, D in self.named_children():
            out = D(input)
            if not get_intermediate_features:
                out = [out]
            result.append(out)
            input = self.downsample(input)
        return result

class NLayerDiscriminator(BaseNetwork):
    def __init__(self, opt):
        super(NLayerDiscriminator, self).__init__()
        self.opt = opt
        kw = 2
        padw = int((kw - 1.0) / 2)
        nf = opt.ngf
        input_nc = self.compute_D_input_nc(opt)
        norm_layer = get_nonspade_norm_layer(norm_type='spectralinstance')
        sequence = [[nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw),
                     nn.LeakyReLU(0.2, False)]]

        for n in range(1, opt.n_layers_D):
            ngf_pre = nf
            nf = min(nf * 2, 512)
            stride = 1 if n == opt.n_layers_D - 1 else 2  # only one layer stride = 1
            if n == opt.n_layers_D - 1:
                dec = []
                nc_dec = ngf_pre
                for _ in range(opt.n_layers_D - 1):
                    dec += [nn.Upsample(scale_factor=2),
                            norm_layer(nn.Conv2d(nc_dec, int(nc_dec // 2), kernel_size=3, stride=1, padding=1)),
                            nn.LeakyReLU(0.2, False)]
                    nc_dec = int(nc_dec // 2)
                dec += [nn.Conv2d(nc_dec, 1, kernel_size=3, stride=1, padding=1)]
                self.dec = nn.Sequential(*dec)
            sequence += [[norm_layer(nn.Conv2d(ngf_pre, nf, kernel_size=kw,
                                               stride=stride, padding=padw)),
                          nn.LeakyReLU(0.2, False)
                          ]]
        sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
        # We divide the layers into groups to extract intermediate layer outputs
        for n in range(len(sequence)):
            self.add_module('model' + str(n), nn.Sequential(*sequence[n]))

    def forward(self, input):
        results = [input]
        for name, submodel in self.named_children():
            if 'model' not in name:
                continue
            x = results[-1]
            intermediate_output = submodel(x)
            results.append(intermediate_output)

        get_intermediate_features = not self.opt.no_ganFeat_loss
        if get_intermediate_features:
            retu = results[1:]
        else:
            retu = results[-1]
        return retu

    def compute_D_input_nc(self, opt):
        if opt.D_type == 'structure':
            input_nc = 4
        elif opt.D_type == 'color':
            input_nc = 6
        else:
            raise ValueError('can not found value'+opt.D_type)
        return input_nc

if __name__ == '__main__':
    opt = Options().parse()
    bs = 2
    sketch = torch.randn(size=(bs,1,256,256))
    ref = torch.randn(size=(bs,3,256,256))
    D = MultiscaleDiscriminator(opt)
    input = torch.cat([sketch,ref], dim=1)
    out = D(input)
    print(out[0][-1].shape)